from transformers import AutoTokenizer, AutoModelForCausalLM,AutoConfig
import torch


config = AutoConfig.from_pretrained("a686d380/rwkv-5-h-world",trust_remote_code=True)  #hugging face对应名称

# 创建一个空的模型实例
model = AutoModelForCausalLM.from_config(config,trust_remote_code=True)

# 加载 .pth 模型的权重
pth_model_path = r"D:\rwkv\models\rwkv-5-h-world-7B.pth"

pth_weights = torch.load(pth_model_path,map_location='cuda:0')

# 将 .pth 模型的权重加载到新模型中
# 这里需要确保 .pth 模型的权重名称与新模型的权重名称匹配。调整键名（示例：假设 .pth 权重键名前缀为 "old_prefix."，需去掉）
adjusted_weights = {k.replace("old_prefix.", ""): v for k, v in pth_weights.items()}
# 可能需要进行一些名称映射的调整
model.load_state_dict(adjusted_weights)

# 保存模型为 transformers 格式
save_dir = r"D:\rwkv\models\rwkv5h1"
print("Before saving model...")
model.save_pretrained(save_dir)
print("After saving model...")
#保存分词器（如果需要）
# 加载分词器
# tokenizer = AutoTokenizer.from_pretrained("gpt2")
# 保存分词器
# tokenizer.save_pretrained(save_dir)